[ExecuTorch][WebGPU] Add softmax ops (_softmax, _log_softmax) to the WebGPU backend#20919
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20919
Note: Links to docs will display an error until the docs builds have been completed. ❌ 53 New Failures, 6 Unclassified FailuresAs of commit 387daf1 with merge base 21554e5 ( NEW FAILURES - The following jobs have failed:
UNCLASSIFIED FAILURES - DrCI could not classify the following jobs because the workflow did not run on the merge base. The failures may be pre-existing on trunk or introduced by this PR:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
This PR needs a
|
SS-JIA
left a comment
There was a problem hiding this comment.
Review automatically exported from Phabricator review in Meta.
…WebGPU backend Pull Request resolved: #20919 Add `aten._softmax` + `aten._log_softmax` handlers (max-stable online softmax) for the training tail. Key changes: - `runtime/ops/{softmax,log_softmax}/` — max-stable softmax + log_softmax WGSL kernels + handlers - `CMakeLists.txt` `WEBGPU_SRCS` — wire the sources Reuses the shared Vulkan partitioner (`aten._softmax`/`_log_softmax` already registered); WebGPU kernels only. Co-authored-with: Claude Code. ghstack-source-id: 405026047 @exported-using-ghexport Differential Revision: [D111755118](https://our.internmc.facebook.com/intern/diff/D111755118/)
…WebGPU backend Pull Request resolved: #20919 Add `aten._softmax` + `aten._log_softmax` handlers (max-stable online softmax) for the training tail. Key changes: - `runtime/ops/{softmax,log_softmax}/` — max-stable softmax + log_softmax WGSL kernels + handlers - `CMakeLists.txt` `WEBGPU_SRCS` — wire the sources Reuses the shared Vulkan partitioner (`aten._softmax`/`_log_softmax` already registered); WebGPU kernels only. Co-authored-with: Claude Code. ghstack-source-id: 405026047 @exported-using-ghexport Differential Revision: [D111755118](https://our.internmc.facebook.com/intern/diff/D111755118/)
Stack from ghstack (oldest at bottom):
Add
aten._softmax+aten._log_softmaxhandlers (max-stable online softmax) for the training tail.Key changes:
runtime/ops/{softmax,log_softmax}/— max-stable softmax + log_softmax WGSL kernels + handlersCMakeLists.txtWEBGPU_SRCS— wire the sourcesReuses the shared Vulkan partitioner (
aten._softmax/_log_softmaxalready registered); WebGPU kernels only.Co-authored-with: Claude Code.
@exported-using-ghexport
Differential Revision: D111755118
Differential Revision: D111755118